package com.tomato.dynamic.threadpool.core.support;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.*;

/**
 * 线程池配置支持
 * 参考 spring 的配置 {@link ThreadPoolTaskExecutor}.
 * @author lizhifu
 * @date 2022/4/8
 */
@Slf4j
public class DynamicExecutorSupport extends ThreadPoolExecutor implements DisposableBean {
    /**
     * awaitTerminationSeconds 等待时间
     */
    protected int awaitTerminationSeconds = 0;
    /**
     * 是否在关闭线程池时等待计划任务完成
     * 不中断正在运行的任务并执行队列中的所有任务。
     */
    protected boolean waitForTasksToCompleteOnShutdown = false;
    /**
     * 线程池名称
     */
    protected String threadPoolName;

    public DynamicExecutorSupport(int corePoolSize,
                               int maximumPoolSize,
                               long keepAliveTime,
                               TimeUnit unit,
                               BlockingQueue<Runnable> workQueue,
                               ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    /**
     * BeanFactory 销毁
     */
    @Override
    public void destroy(){
        log.info("DynamicExecutorSupport 销毁: {}", threadPoolName);
        internalShutdown();
    }

    public void setAwaitTerminationSeconds(int awaitTerminationSeconds) {
        this.awaitTerminationSeconds = awaitTerminationSeconds;
    }

    /**
     * 关闭线程池
     *
     * @see ExecutorService#shutdown() 该状态下线程池不再接受新任务，但是会将工作队列中的任务执行结束。线程状态：SHUTDOWN
     * @see ExecutorService#shutdownNow() 该状态下线程池不再接受新任务，也不会处理工作队列中的任务，并且将会中断线程。线程状态：STOP
     */
    public void internalShutdown() {
        log.info("关闭线程池: {}", threadPoolName);
        if (this.waitForTasksToCompleteOnShutdown) {
            this.shutdown();
        } else {
            for (Runnable remainingTask : this.shutdownNow()) {
                cancelRemainingTask(remainingTask);
            }
        }
        // 等待线程池终止
        awaitTerminationIfNecessary();
    }
    /**
     * 取消从未执行过的给定剩余任务,
     * as returned from {@link ExecutorService#shutdownNow()}.
     *
     * @param task 要取消的任务 (typically a {@link RunnableFuture})
     * @see #shutdown()
     * @see RunnableFuture#cancel(boolean)
     * @since 5.0.5
     */
    protected void cancelRemainingTask(Runnable task) {
        if (task instanceof Future) {
            ((Future<?>) task).cancel(true);
        }
    }
    /**
     * 线程池已经关闭且内部的任务都已经执行完毕了
     */
    private void awaitTerminationIfNecessary() {
        if (this.awaitTerminationSeconds <= 0) {
            return;
        }
        try {
            if (!awaitTermination(this.awaitTerminationSeconds, TimeUnit.SECONDS) && log.isWarnEnabled()) {
                log.warn("等待线程池终止超时 {}", threadPoolName);
            }
        } catch (InterruptedException ex) {
            if (log.isWarnEnabled()) {
                log.warn("等待线程池终止超时被打断 {}", threadPoolName);
            }
            Thread.currentThread().interrupt();
        }
    }

    public void setThreadPoolName(String threadPoolName) {
        this.threadPoolName = threadPoolName;
    }

    public String getThreadPoolName() {
        return threadPoolName;
    }

    public void setWaitForTasksToCompleteOnShutdown(boolean waitForTasksToCompleteOnShutdown) {
        this.waitForTasksToCompleteOnShutdown = waitForTasksToCompleteOnShutdown;
    }
}
